#语义检索模块

import json
import os
from typing import List, Dict

class SemanticRetriever:
    def __init__(self, kb_path='rust_knowledge_base/rust_docs_sample.json'):
        self.kb_path = kb_path
        self.knowledge_base = self._load_knowledge_base()
    
    def _load_knowledge_base(self) -> List[Dict]:
        """加载知识库"""
        if os.path.exists(self.kb_path):
            with open(self.kb_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        else:
            # 默认知识库
            return [
                {
                    "id": "k001",
                    "topic": "生命周期",
                    "content": "Rust的生命周期用于确保引用始终有效。生命周期是引用保持有效的作用域。",
                    "tags": ["生命周期", "引用", "作用域"],
                    "category": "definition",
                    "code": "fn longest<'a>(x: &'a str, y: &'a str) -> &'a str {\n    if x.len() > y.len() { x } else { y }\n}"
                },
                {
                    "id": "k002", 
                    "topic": "所有权",
                    "content": "Rust的所有权系统是内存安全的核心。每个值都有一个所有者，当所有者离开作用域时，值会被丢弃。",
                    "tags": ["所有权", "内存管理", "作用域"],
                    "category": "definition",
                    "code": "let s1 = String::from(\"hello\");\nlet s2 = s1; // s1的所有权移动到s2"
                },
                {
                    "id": "k003",
                    "topic": "可变变量",
                    "content": "使用mut关键字声明可变变量，允许修改变量的值。",
                    "tags": ["变量", "mut", "可变性"],
                    "category": "usage",
                    "code": "let mut x = 5;\nx = 10; // 可以修改"
                },
                {
                    "id": "k004",
                    "topic": "借用检查器错误",
                    "content": "借用检查器错误通常是由于违反了Rust的借用规则导致的。常见错误包括同时存在可变和不可变借用。",
                    "tags": ["借用检查器", "错误", "调试"],
                    "category": "error_debug",
                    "code": "let mut v = vec![1, 2, 3];\nlet first = &v[0];\nv.push(4); // 错误：同时存在可变和不可变借用"
                },
                {
                    "id": "k005",
                    "topic": "Rust vs C++",
                    "content": "Rust和C++都是系统编程语言，但Rust通过所有权系统提供内存安全，而C++需要手动管理内存。",
                    "tags": ["Rust", "C++", "比较"],
                    "category": "comparison",
                    "code": "// Rust: 自动内存管理\nlet s = String::from(\"hello\");\n// C++: 手动内存管理\n// std::string s = \"hello\";"
                }
            ]
    
    def retrieve(self, question: str, question_type: str = None, top_k: int = 3) -> List[Dict]:
        """语义检索相关文档"""
        # 简单的关键词匹配 + 类型过滤
        results = []
        
        for doc in self.knowledge_base:
            score = 0
            
            # 基于问题类型的过滤
            if question_type and doc.get('category') == question_type:
                score += 0.5
            
            # 关键词匹配
            question_words = set(question.lower().split())
            content_words = set(doc['content'].lower().split())
            tag_words = set([tag.lower() for tag in doc.get('tags', [])])
            
            # 计算匹配度
            content_match = len(question_words & content_words) / len(question_words) if question_words else 0
            tag_match = len(question_words & tag_words) / len(question_words) if question_words else 0
            
            score += content_match * 0.3 + tag_match * 0.2
            
            if score > 0:
                results.append({
                    **doc,
                    'score': score
                })
        
        # 按分数排序并返回top_k
        results.sort(key=lambda x: x['score'], reverse=True)
        return results[:top_k]
    
    def add_document(self, doc: Dict):
        """添加新文档到知识库"""
        if 'id' not in doc:
            doc['id'] = f"k{len(self.knowledge_base) + 1:03d}"
        self.knowledge_base.append(doc)
        
        # 保存到文件
        with open(self.kb_path, 'w', encoding='utf-8') as f:
            json.dump(self.knowledge_base, f, ensure_ascii=False, indent=2) 